package data_store

import (
	"sync"

	"atomgit.com/openlinksaas/proto-gen-go.git/trace_proxy_api"
	"github.com/edwingeng/deque/v2"
)

type TraceRecord struct {
	TraceId    string
	RootSpanId string
}

func (record *TraceRecord) ToTraceInfo(rootSpan *trace_proxy_api.SpanInfo) *trace_proxy_api.TraceInfo {
	return &trace_proxy_api.TraceInfo{
		TraceId:  record.TraceId,
		RootSpan: rootSpan,
	}
}

type TraceTable struct {
	traceMap   map[string]*TraceRecord
	traceQueue *deque.Deque[string]
	mutex      sync.Mutex

	//相关配置
	maxStore uint32 //保存最大条数
}

func newTraceTable(maxStore uint32) *TraceTable {
	return &TraceTable{
		traceMap:   make(map[string]*TraceRecord, maxStore),
		traceQueue: deque.NewDeque[string](),
		maxStore:   maxStore,
	}
}

func (tab *TraceTable) Insert(record *TraceRecord) {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	_, ok := tab.traceMap[record.TraceId]
	if ok {
		return
	}

	tab.traceMap[record.TraceId] = record

	tab.traceQueue.PushFront(record.TraceId)
}

func (tab *TraceTable) UpdateRootSpanId(traceId, rootSpanId string) {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	record, ok := tab.traceMap[traceId]
	if ok {
		record.RootSpanId = rootSpanId
	}
}

func (tab *TraceTable) RemoveOldTrace() string {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	if len(tab.traceMap) < int(tab.maxStore) {
		return ""
	}
	traceId := tab.traceQueue.PopBack()
	delete(tab.traceMap, traceId)

	return traceId
}

func (tab *TraceTable) Get(traceId string) (*TraceRecord, bool) {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	record, ok := tab.traceMap[traceId]
	return record, ok
}

func (tab *TraceTable) List() []*TraceRecord {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	retList := []*TraceRecord{}

	for _, record := range tab.traceMap {
		if record.RootSpanId != "" {
			retList = append(retList, record)
		}
	}

	return retList
}
